// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/plan/plan_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>

using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::aggregate::test {

namespace {

class BitwiseAggregationTest : public AggregationTestBase {
 protected:
  void SetUp() override {
    AggregationTestBase::SetUp();
  }

  RowTypePtr rowType_{
      ROW({"c0", "c1", "c2", "c3", "c4"},
          {BIGINT(), TINYINT(), SMALLINT(), INTEGER(), BIGINT()})};
};

TEST_F(BitwiseAggregationTest, bitwiseOr) {
  auto vectors = makeVectors(rowType_, 10, 100);
  createDuckDbTable(vectors);

  // Global aggregation.
  testAggregations(
      vectors,
      {},
      {"bitwise_or_agg(c1)",
       "bitwise_or_agg(c2)",
       "bitwise_or_agg(c3)",
       "bitwise_or_agg(c4)"},
      "SELECT bit_or(c1), bit_or(c2), bit_or(c3), bit_or(c4) FROM tmp");

  // Group by partial aggregation.
  testAggregations(
      [&](auto& builder) {
        builder.values(vectors).project({"c0 % 10", "c1", "c2", "c3", "c4"});
      },
      {"p0"},
      {"bitwise_or_agg(c1)",
       "bitwise_or_agg(c2)",
       "bitwise_or_agg(c3)",
       "bitwise_or_agg(c4)"},
      "SELECT c0 % 10, bit_or(c1), bit_or(c2), bit_or(c3), bit_or(c4) FROM tmp GROUP BY 1");
}

TEST_F(BitwiseAggregationTest, bitwiseAnd) {
  auto vectors = makeVectors(rowType_, 10, 100);
  createDuckDbTable(vectors);

  // Global aggregation.
  testAggregations(
      vectors,
      {},
      {"bitwise_and_agg(c1)",
       "bitwise_and_agg(c2)",
       "bitwise_and_agg(c3)",
       "bitwise_and_agg(c4)"},
      "SELECT bit_and(c1), bit_and(c2), bit_and(c3), bit_and(c4) FROM tmp");

  // Group by aggregation.
  testAggregations(
      [&](auto& builder) {
        builder.values(vectors).project({"c0 % 10", "c1", "c2", "c3", "c4"});
      },
      {"p0"},
      {"bitwise_and_agg(c1)",
       "bitwise_and_agg(c2)",
       "bitwise_and_agg(c3)",
       "bitwise_and_agg(c4)"},
      "SELECT c0 % 10, bit_and(c1), bit_and(c2), bit_and(c3), bit_and(c4) FROM tmp GROUP BY 1");
}

TEST_F(BitwiseAggregationTest, bitwiseXor) {
  auto vectors = makeVectors(rowType_, 10, 100);
  createDuckDbTable(vectors);

  // Global aggregation.
  testAggregations(
      vectors,
      {},
      {"bitwise_xor_agg(c1)",
       "bitwise_xor_agg(c2)",
       "bitwise_xor_agg(c3)",
       "bitwise_xor_agg(c4)"},
      "SELECT bit_xor(c1), bit_xor(c2), bit_xor(c3), bit_xor(c4) FROM tmp");

  // Group by aggregation.
  testAggregations(
      [&](auto& builder) {
        builder.values(vectors).project({"c0 % 10", "c1", "c2", "c3", "c4"});
      },
      {"p0"},
      {"bitwise_xor_agg(c1)",
       "bitwise_xor_agg(c2)",
       "bitwise_xor_agg(c3)",
       "bitwise_xor_agg(c4)"},
      "SELECT c0 % 10, bit_xor(c1), bit_xor(c2), bit_xor(c3), bit_xor(c4) FROM tmp GROUP BY 1");
}

} // namespace
} // namespace kumo::pollux::aggregate::test
